#pragma once
#include <atomic>
#include <condition_variable>
#include <map>
#include <memory>
#include <mutex>
#include <string>
#include <unordered_map>
#include <vector>

#include "ring_buffer.h"
#include "thread.h"
#define DISABLE_COPY_AND_ASSIGN(ClassType) \
  ClassType(const ClassType &) = delete;   \
  ClassType &operator=(const ClassType &) = delete;

constexpr uint32_t kDefaultRingBuffer = 1024;
constexpr uint32_t kBatchMaxLen =  1024;
constexpr uint32_t kMaxWaitTimeUs = 1000;
constexpr uint32_t kNotifyInterval = 2 * 256;

enum class OpRangeDataType : uint8_t {
  THREAD_ID = 0,
  FLOW_ID = 1,
  STEP = 2,
  START_TIME_NS = 3,
  END_TIME_NS = 4,
  PROCESS_ID = 5,
  MODULE_INDEX = 6,
  EVENT_INDEX = 7,
  STAGE_INDEX = 8,
  LEVEL = 9,
  IS_ASYNC = 10,
  IS_GRAPH_DATA = 11,
  IS_STAGE = 12,
  NAME = 13,
};

enum class ReportFileType : uint32_t {
  OP_RANGE = 0,
  OP_MARK = 1,
  PYTHON_STACK = 2,
};

static const std::map<ReportFileType, std::string> kReportFileTypeMap = {
    {ReportFileType::OP_RANGE, "mindspore.op_range"},
    {ReportFileType::OP_MARK, "mindspore.op_mark"},
    {ReportFileType::PYTHON_STACK, "mindspore.py_stack"},
};

struct BaseReportData {
  // 1 * 4 bytes = 4 bytes
  int32_t device_id{0};
  // 1 * 4 bytes = 4 bytes
  uint32_t tag;
  BaseReportData(int32_t device_id, uint32_t tag)
      : device_id(device_id), tag(tag) {}
  virtual ~BaseReportData() = default;
  virtual std::vector<uint8_t> encode() = 0;
};

struct MindSporeOpRangeData : BaseReportData {
  // 5 * 8 bytes = 40 bytes
  uint64_t thread_id{0};
  uint64_t flow_id{0};
  uint64_t step{0};
  uint64_t start_time_ns{0};
  uint64_t end_time_ns{0};

  // 1 * 4 bytes = 4 bytes
  int32_t process_id{0};

  // 3 * 2 bytes = 6 bytes
  uint16_t module_index{0};
  uint16_t event_index{0};
  uint16_t stage_index{0};

  // 4 * 1 bytes = 4 bytes
  int8_t level{-1};
  bool is_async{false};
  bool is_graph_data{false};
  bool is_stage{false};

  // dynamic length
  std::string op_name;

  MindSporeOpRangeData(int32_t device_id, uint32_t tag, uint64_t thread_id,
                       uint64_t flow_id, uint64_t step, uint64_t start_time_ns,
                       uint64_t end_time_ns, int32_t process_id,
                       uint16_t module_index, uint16_t event_index,
                       uint16_t stage_index, int8_t level, bool is_async,
                       bool is_graph_data, bool is_stage, std::string op_name)
      : BaseReportData(device_id,
                       static_cast<uint32_t>(ReportFileType::OP_RANGE)),
        thread_id(thread_id),
        flow_id(flow_id),
        step(step),
        start_time_ns(start_time_ns),
        end_time_ns(end_time_ns),
        process_id(process_id),
        module_index(module_index),
        event_index(event_index),
        stage_index(stage_index),
        level(level),
        is_async(is_async),
        is_graph_data(is_graph_data),
        is_stage(is_stage),
        op_name(std::move(op_name)) {}

  constexpr static size_t tlvPrefixLength = sizeof(uint16_t) + sizeof(uint32_t);
  constexpr static size_t tlvFixedDataLength =
      sizeof(uint64_t) * 5 + sizeof(int32_t) + sizeof(uint16_t) * 3 +
      sizeof(int8_t) + sizeof(bool) * 3;

  // TLV buffer length = Type + Length + Fixed data
  constexpr static size_t tlvBufferLength = tlvPrefixLength + tlvFixedDataLength;
  std::vector<uint8_t> encode() override;
};

template <typename T>
inline void encodeFixedData(const T &data, std::vector<uint8_t> &result) {
  for (size_t i = 0; i < sizeof(T); ++i) {
    result.emplace_back((static_cast<size_t>(data) >> (i << 3)) & 0xff);
  }
}


inline void encodeStrData(uint16_t type, const std::string &data,
                          std::vector<uint8_t> &result) {
  for (size_t i = 0; i < sizeof(uint16_t); ++i) {
    result.emplace_back((type >> (i << 3)) & 0xff);
  }
  uint32_t length = data.size();
  for (size_t i = 0; i < sizeof(uint32_t); ++i) {
    result.emplace_back((length >> (i << 3)) & 0xff);
  }
  for (const auto &c : data) {
    result.emplace_back(c);
  }
}

class DataDumper : public Thread {
 public:
  explicit DataDumper();
  virtual ~DataDumper();
  void Init(const std::string &path, size_t capacity);
  void UnInit();
  void Report(std::unique_ptr<BaseReportData> data);
  void Start();
  void Stop();
  static DataDumper &GetInstance();

 private:
  void Flush();
  void Dump(
      const std::unordered_map<std::string, std::vector<uint8_t>> &dataMap);
  void Run();
  void GatherAndDumpData();

 private:
  std::string path_;
  std::atomic<bool> start_;
  std::atomic<bool> init_;
  RingBuffer<std::unique_ptr<BaseReportData>> data_chunk_buf_;
  std::unordered_map<std::string, FILE *> fd_map_;
  uint64_t dump_count_{0};
};
